'''
https://leetcode.cn/problems/shortest-word-distance-ii/
https://leetcode.cn/problems/shortest-word-distance-iii
'''

from typing import List
from collections import defaultdict

class WordDistance:

    def __init__(self, wordsDict: List[str]):
        self.word_indices = word_indices = defaultdict(list)
        for i, word in enumerate(wordsDict):
            word_indices[word].append(i)

    def shortest(self, word1: str, word2: str) -> int:
        i, j = 0, 0
        l1, l2 = self.word_indices[word1], self.word_indices[word2]
        res = 100000000
        while i < len(l1) and j < len(l2):
            idx1 = l1[i]
            idx2 = l2[j]
            res = min(res, abs(idx1 - idx2))
            if idx1 < idx2:
                i += 1
            else:
                j += 1
        return res


class Solution:
    def shortestWordDistance(self, wordsDict: List[str], word1: str, word2: str) -> int:
        l1, l2 = [],[]
        for i, word in enumerate(wordsDict):
            if word == word1:
                l1.append(i)
            elif word == word2:
                l2.append(i)

        if not l2:
            return min(abs(l1[i] - l1[i-1]) for i in range(1, len(l1)))
        i, j = 0, 0
        res = 100000000
        while i < len(l1) and j < len(l2):
            idx1 = l1[i]
            idx2 = l2[j]
            res = min(res, abs(idx1 - idx2))
            if idx1 < idx2:
                i += 1
            else:
                j += 1
        return res
    
wordsDict = ["practice", "makes", "perfect", "coding", "makes"]
word1 = 'makes'
word2 = 'coding'
print(Solution().shortestWordDistance(wordsDict, word1, word2))